from datasets import load_dataset
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.svm import LinearSVC
from sklearn.metrics import accuracy_score
from joblib import dump, load
import os
import time

# 控制是否跳过训练并直接加载已训练好的模型
skip_training = False  # 设置为 True 跳过训练，False 重新训练

# 配置模型路径和数据路径
trained_model_path = "./results/svm_sentiment_model.joblib"  # 保存最终模型的路径
data_files = {
    "train": "./imdb/plain_text/train-00000-of-00001.parquet",
    "test": "./imdb/plain_text/test-00000-of-00001.parquet",
}

# 加载数据集
print("Loading dataset...")
dataset = load_dataset("parquet", data_files=data_files)

# 提取训练集和测试集文本与标签
train_texts = [x["text"] for x in dataset["train"]]
train_labels = [x["label"] for x in dataset["train"]]
test_texts = [x["text"] for x in dataset["test"]]
test_labels = [x["label"] for x in dataset["test"]]

if not skip_training:
    print("Training model...")

    # 记录训练开始时间
    time_start = time.time()

    # 1. 特征提取：使用词袋模型将文本转换为特征向量
    vectorizer = CountVectorizer(max_features=5000, stop_words="english")  # 使用 5000 个词作为特征
    train_features = vectorizer.fit_transform(train_texts)
    test_features = vectorizer.transform(test_texts)

    # 2. 训练支持向量机 (SVM) 分类器
    classifier = LinearSVC()
    classifier.fit(train_features, train_labels)

    # 记录训练结束时间
    time_end = time.time()

    # 保存最终模型和词袋模型
    os.makedirs(os.path.dirname(trained_model_path), exist_ok=True)
    dump((classifier, vectorizer), trained_model_path)
    print(f"Model saved to {trained_model_path}")

    # 打印训练时间
    print(f"Training time: {time_end - time_start:.2f} seconds")

else:
    print("Loading pre-trained model...")
    # 加载已保存的模型
    classifier, vectorizer = load(trained_model_path)

# 测试模型性能
print("Evaluating model...")
test_features = vectorizer.transform(test_texts)
predictions = classifier.predict(test_features)
accuracy = accuracy_score(test_labels, predictions)
print(f"Final Test Accuracy: {accuracy:.4f}")

# 测试单条输入文本的情感分类
text = "I absolutely loved this movie! It was amazing and so well acted."
text_feature = vectorizer.transform([text])
predicted_class = classifier.predict(text_feature)[0]
print("Predicted Sentiment:", "Positive" if predicted_class == 1 else "Negative")
